#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <elf.h>
#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include "symvers.h"

#if 1
#define DBG_PRINT       printf
#else
#define DBG_PRINT
#endif

#if __WORDSIZE == 64
#define ELFCLASS ELFCLASS64
typedef Elf64_Ehdr Elf_Ehdr;
typedef Elf64_Phdr Elf_Phdr;
typedef Elf64_Shdr Elf_Shdr;
typedef Elf64_Off  Elf_Off;
#else
#define ELFCLASS ELFCLASS32
typedef Elf32_Ehdr Elf_Ehdr;
typedef Elf32_Phdr Elf_Phdr;
typedef Elf32_Shdr Elf_Shdr;
typedef Elf32_Off  Elf_Off;
#endif

struct kernel_symbol
{
	unsigned long value;//address of symbol
	const char *name;//name of symbol
};

static struct symvers_info g_symvers = {0};

#define SECTION_KCRCTAB_STR             "__kcrctab"
#define SECTION_KCRCTAB_GPL_STR         "__kcrctab_gpl"
#define SECTION_KSYMTAB_STR             "__ksymtab"
#define SECTION_KSYMTAB_GPL_STR         "__ksymtab_gpl"
#define SECTION_KSYMTAB_STRINGS_STR     "__ksymtab_strings"

static char *inner_find_section(Elf_Ehdr *ehdr, const char *section_name, unsigned int *pSectionLen, Elf_Shdr **pshdr)
{
    char *p = NULL;
    Elf_Shdr *shdr = (Elf_Shdr *)((unsigned long)ehdr + ehdr->e_shoff);
    Elf_Shdr *shstrtab = &shdr[ehdr->e_shstrndx];
    unsigned long shstrtab_off = (unsigned int)shstrtab->sh_offset;
    char *strtab = (char *)((unsigned long)ehdr + shstrtab_off);

    *pSectionLen = 0;
    for (int i = 0; i < (int)ehdr->e_shnum; i++)
    {
        if (0 == strcmp(strtab + shdr[i].sh_name, section_name))
        {
            p = (char *)((unsigned long)ehdr + shdr[i].sh_offset);
            *pSectionLen = shdr[i].sh_size;
            if(NULL != pshdr)
            {
                *pshdr = &(shdr[i]);
            }
            break;
        }
    }
    return p;
}

static char *inner_strcpy_s(char *buffer, unsigned int buflen, const char *src)
{
    char *dest = buffer;

    while('\0' != *src && buflen > 1)
    {
        *dest = *src;
        ++src;
        ++dest;
        --buflen;
    }
    if(buflen >= 1)
    {
        (*dest) = '\0';
    }
    return buffer;
}

static void inner_parse_symvers(unsigned long base_address, const unsigned long *crctab, const struct kernel_symbol *symtab,
    Elf_Off symtab_strings_offset , unsigned int symtab_strings_max_offset, unsigned char *map_addr,
    unsigned int count, struct sym_info *list)
{
    unsigned long offset;
    const char *sym_name;

    for(unsigned int i = 0; i < count; i++)
    {
        offset = (unsigned long)symtab[i].name - base_address;
        if(offset < symtab_strings_max_offset && offset >= symtab_strings_offset)
        {
            list[i].crc = crctab[i];
            sym_name = (char *)(map_addr + offset);
            inner_strcpy_s(list[i].name, sizeof(list[i].name), sym_name);
        }
    }
}

static int inner_parse_vmlinux(const char *path)
{
    int ret = 0;
    int fd;
    struct stat st;
    unsigned char *addr = (unsigned char *)MAP_FAILED;

    fd = open(path, O_RDONLY);
    if (-1 == fd)
    {
        ret = errno;
        DBG_PRINT("open <%s> failed!ret=%d\n", path, ret);
        goto _OUT;
    }
    if (0 != fstat(fd, &st))
    {
        ret = errno;
        DBG_PRINT("fstat <%s> failed!ret=%d\n", path, ret);
        goto _OUT;
    }
    addr = (unsigned char *)mmap(NULL, st.st_size, PROT_READ, MAP_SHARED, fd, 0);
    if (MAP_FAILED == addr)
    {
        ret = errno;
        DBG_PRINT("mmap <%s> failed!ret=%d\n", path, ret);
        goto _OUT;
    }

    if (0 != memcmp(addr, ELFMAG, SELFMAG))
    {
        DBG_PRINT("Invalid ELF file<%s>!\n", path);
        ret = EINVAL;
        goto _OUT;
    }

    if (ELFCLASS == addr[4])
    {
        Elf_Ehdr *ehdr = (Elf_Ehdr *)addr;
        Elf_Shdr *shdr = NULL;
        const unsigned long *crctab;
        const unsigned long *crctab_gpl;
        const struct kernel_symbol *symtab;
        const struct kernel_symbol *symtab_gpl;
        const char *symtab_strings;
        unsigned int crctab_len = 0;
        unsigned int crctab_gpl_len = 0;
        unsigned int symtab_len = 0;
        unsigned int symtab_gpl_len = 0;
        unsigned int symtab_strings_len = 0;
        unsigned long base_address;

        symtab_strings = (const char *)inner_find_section(ehdr, SECTION_KSYMTAB_STRINGS_STR, &symtab_strings_len, &shdr);
        if(NULL == symtab_strings || 0 == symtab_strings_len || NULL == shdr)
        {
            DBG_PRINT("Don't find section<%s> in %s!\n", SECTION_KSYMTAB_STRINGS_STR, path);
            ret = ENOENT;
            goto _OUT;
        }
        crctab = (const unsigned long *)inner_find_section(ehdr, SECTION_KCRCTAB_STR, &crctab_len, NULL);
        if(NULL == crctab || 0 == crctab_len)
        {
            DBG_PRINT("Don't find section<%s> in %s!\n", SECTION_KCRCTAB_STR, path);
            ret = ENOENT;
            goto _OUT;
        }
        symtab = (const struct kernel_symbol *)inner_find_section(ehdr, SECTION_KSYMTAB_STR, &symtab_len, NULL);
        if(NULL == symtab || 0 == symtab_len)
        {
            DBG_PRINT("Don't find section<%s> in %s!\n", SECTION_KSYMTAB_STR, path);
            ret = ENOENT;
            goto _OUT;
        }
        crctab_gpl = (const unsigned long *)inner_find_section(ehdr, SECTION_KCRCTAB_GPL_STR, &crctab_gpl_len, NULL);
        symtab_gpl = (const struct kernel_symbol *)inner_find_section(ehdr, SECTION_KSYMTAB_GPL_STR, &symtab_gpl_len, NULL);

        g_symvers.sym_count = crctab_len/sizeof(unsigned long);
        if (crctab_gpl_len > 0 && NULL != crctab_gpl && symtab_gpl_len > 0 && NULL != symtab_gpl)
        {
            g_symvers.symgpl_count = crctab_gpl_len / sizeof(unsigned long);
        }
        if(g_symvers.sym_count > 0 || g_symvers.symgpl_count > 0)
        {
            g_symvers.list = (struct sym_info *)malloc((g_symvers.sym_count + g_symvers.symgpl_count) * sizeof(struct sym_info));
            if(NULL == g_symvers.list)
            {
                DBG_PRINT("malloc failed!\n");
                ret = ENOMEM;
                g_symvers.sym_count = 0;
                g_symvers.symgpl_count = 0;
                goto _OUT;
            }
            base_address = shdr->sh_addr - shdr->sh_offset;
            if(g_symvers.sym_count > 0)
            {
                inner_parse_symvers(base_address, crctab, symtab, shdr->sh_offset, symtab_strings_len + shdr->sh_offset,
                    addr, g_symvers.sym_count, g_symvers.list);
            }
            if(g_symvers.symgpl_count > 0)
            {
                inner_parse_symvers(base_address, crctab_gpl, symtab_gpl, shdr->sh_offset, symtab_strings_len + shdr->sh_offset,
                    addr, g_symvers.symgpl_count, g_symvers.list + g_symvers.sym_count);
            }
        }
    }
    else
    {
        ret = EINVAL;
    }
_OUT:
    if (MAP_FAILED != addr)
    {
        munmap(addr, st.st_size);
    }
    if (-1 != fd)
    {
        close(fd);
    }
    return ret;
}

static int inner_extract_vmlinuz(const char *path, const char *target)
{
    char cmdbuf[2048] = {0};

    snprintf(cmdbuf, sizeof(cmdbuf) - 1, "/usr/bin/sh extract-vmlinux \"%s\" \"%s\"", path, target);

    remove(target);
    system(cmdbuf);
    if(0 == access(target, F_OK))
    {
        return 0;
    }
    return ENOENT;
}

int symvers_init(const char *path)
{
    int ret;
    // /tmp/.vmz.d0cd2c8a-631e-4d5f-b8fe-6af6c5ec28f8
    char target[] = {'/', 't', 'm', 'p', '/', '.', 'v', 'm', 'z', '.', 'd', '0', 'c', 'd', '2', 'c', '8',
                     'a', '-', '6', '3', '1', 'e', '-', '4', 'd', '5', 'f', '-', 'b', '8', 'f', 'e',
                     '-', '6', 'a', 'f', '6', 'c', '5', 'e', 'c', '2', '8', 'f', '8', '\0'};

    if(NULL == path || '\0' == *path)
    {
        return EINVAL;
    }
    if(NULL != g_symvers.list)
    {
        return EEXIST;
    }
    ret = inner_extract_vmlinuz(path, target);
    if(0 != ret)
    {
        return ret;
    }
    ret = inner_parse_vmlinux(target);
    remove(target);
    return ret;
}

void symvers_uninit(void)
{
    if(NULL != g_symvers.list)
    {
        free(g_symvers.list);
        memset(&g_symvers, 0, sizeof(g_symvers));
    }
}

const struct symvers_info *symvers_get_verinfo_list(void)
{
    return &g_symvers;
}
